import numpy as np
from matplotlib import pyplot as plt
from sklearn.linear_model import LinearRegression

"""构造数据"""
x = 2 * np.random.rand(100, 1)
y = 4 + 3 * x + np.random.randn(100, 1)
x_b = np.c_[np.ones([100, 1]), x]  # 加一列

"""做个图看一眼"""
plt.scatter(x, y, label="Train data")
plt.xlabel("x")
plt.ylabel("y")
plt.title("train")
plt.legend()  # 显示label
plt.show()


"""算法实例化"""
lin_reg = LinearRegression()

"""进行数据初始化"""
lin_reg.fit(x, y)

"""梯度下降"""
eta = 0.1  # 学习率
n_epochs = 50  # 迭代次数
mini_batch = 16
m = 100  # 样本个数
theta = np.random.randn(2, 1)  # 参数的随机初始化
t = 0
t0 = 5
t1 = 50
theta_path_mgd = []

def learning_schedule(t):
    return t0/(t1+t)


"""迭代"""
for i in range(n_epochs):
    """先对数据进行洗牌"""
    shuffled_indices = np.random.permutation(m)  # 对一维数组[0-m]进行重排
    x_b_shuffled = x_b[shuffled_indices]
    y_shuffled = y[shuffled_indices]
    for j in range(0, m, mini_batch):
        t += 1
        xi = x_b_shuffled[i:i + mini_batch]  # 每次选择一批
        yi = y_shuffled[i:i + mini_batch]
        gradients = 2 / mini_batch * xi.T.dot(xi.dot(theta) - yi)
        eta = learning_schedule(t)  # 更新学习率
        theta = theta - eta * gradients  # 更新参数列表
        theta_path_mgd.append(theta)

"""测试参数"""
x_new = np.array([[0], [2]])  # 构造新数据
x_new_b = np.c_[np.ones([2, 1]), x_new]  # 加一列偏置项
y_predict = x_new_b.dot(theta)
print(theta)

"""画出来回归方程"""
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
plt.scatter(x, y, label="Train data")
plt.scatter(x_new, y_predict, label="Test data")
plt.plot(x_new, y_predict, 'r', label="回归方程")
plt.xlabel("x")
plt.ylabel("y")
plt.title("Happy")
plt.legend()
plt.show()

